#include <iostream>
#include <fstream>
#include <string>
#include <map>
#include <conio.h> 

#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <winsock2.h>  //使用ip地址相关函数
#include <windows.h>

typedef unsigned char u_char;
typedef unsigned short u_short;
typedef unsigned int u_int;
typedef unsigned long u_long;

#endif

#include <pcap.h>  //引入 WinPcap 的核心头文件



std::ofstream log_file;  // 用于将每个数据包的详细信息写入 traffic.csv
//src_stats / dst_stats：分别统计源地址和目标地址的总流量（按 <MAC, IP> 组合统计）
std::map<std::pair<std::string, std::string>, uint64_t> src_stats;
std::map<std::pair<std::string, std::string>, uint64_t> dst_stats;
time_t last_stat_time;  //记录上一次统计输出的时间戳，用于判断是否需要输出新一轮的统计信息。

//报文处理
void process_packet(const struct pcap_pkthdr* header, const u_char* packet) {
    if (header->caplen < 14) return;  //报文长度不足14字节，不是有效的以太网帧，直接跳过

    u_char dst_mac[6], src_mac[6];  //以太网帧前6字节是 目的MAC，接下来的6字节是 源MAC
    memcpy(dst_mac, packet, 6);
    memcpy(src_mac, packet + 6, 6);  //提取mac地址

    std::cerr << "get MAC successfully" << std::endl;
    //格式化成标准 MAC 地址字符串，形如 AA-BB-CC-DD-EE-FF
    char src_mac_str[18], dst_mac_str[18];
    snprintf(src_mac_str, 18, "%02X-%02X-%02X-%02X-%02X-%02X",
        src_mac[0], src_mac[1], src_mac[2], src_mac[3], src_mac[4], src_mac[5]);
    snprintf(dst_mac_str, 18, "%02X-%02X-%02X-%02X-%02X-%02X",
        dst_mac[0], dst_mac[1], dst_mac[2], dst_mac[3], dst_mac[4], dst_mac[5]);
    //第13-14字节是类型字段，如果不是 0x0800（即 IPv4），则跳过。
    uint16_t eth_type = ntohs(*reinterpret_cast<const uint16_t*>(packet + 12));
    if (eth_type != 0x0800) return;

    std::cerr << "it is IPv4" << std::endl;

    if (header->caplen < 34) return;
    const u_char* ip_start = packet + 14;   //IP 报文从第15字节开始。
    uint8_t ihl = (*ip_start & 0x0F) * 4;   //ihl 表示 IP 首部长度（以 4 字节为单位），用于处理带选项的 IP 报头
    if (ihl < 20 || header->caplen < 14 + ihl) return;

    struct in_addr src_addr, dst_addr;
    //IP 报文的第13-16字节为源 IP，第17-20字节为目标 IP
    memcpy(&src_addr, ip_start + 12, 4);
    memcpy(&dst_addr, ip_start + 16, 4);
    char src_ip[INET_ADDRSTRLEN], dst_ip[INET_ADDRSTRLEN];
    //将 IP 地址从二进制格式转换为点分十进制字符串
    inet_ntop(AF_INET, &src_addr, src_ip, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &dst_addr, dst_ip, INET_ADDRSTRLEN);
    //提取数据包到达的时间，并格式化成人类可读格式
    time_t raw_time = header->ts.tv_sec;
    struct tm* timeinfo = new tm();
    localtime_s(timeinfo, &raw_time);
    char time_str[20];
    strftime(time_str, 20, "%Y-%m-%d %H:%M:%S", timeinfo);
    //输出字段包括时间、源/目标 MAC、IP、帧长度，实时写入 traffic.csv
    log_file << time_str << "," << src_mac_str << "," << src_ip << ","
        << dst_mac_str << "," << dst_ip << "," << header->len << "\n";
    log_file.flush();
    //以 <MAC, IP> 为键，统计对应主机累计发送/接收的流量（字节数）
    src_stats[{src_mac_str, src_ip}] += header->len;
    dst_stats[{dst_mac_str, dst_ip}] += header->len;
}
//输出统计信息
void output_stats() {
    std::ofstream stat_file("statistics.csv", std::ios::app);
    time_t now = time(nullptr);
    struct tm* timeinfo = new tm();
    localtime_s(timeinfo, &now);
    char time_str[20];
    strftime(time_str, 20, "%Y-%m-%d %H:%M:%S", timeinfo);

    for (const auto& entry : src_stats) {
        stat_file << time_str << ",SRC," << entry.first.first << ","
            << entry.first.second << "," << entry.second << "\n";
    }
    for (const auto& entry : dst_stats) {
        stat_file << time_str << ",DST," << entry.first.first << ","
            << entry.first.second << "," << entry.second << "\n";
    }

    src_stats.clear();
    dst_stats.clear();
}

int main() {

    char errbuf[PCAP_ERRBUF_SIZE];   //定义存储错误信息的缓冲区
    //查找并列出所有可用的网络接口，用于选择抓包设备
    pcap_if_t* all_devs;  
    pcap_findalldevs(&all_devs, errbuf);  //查找本机所有可抓包的网卡（返回链表）
    //遍历每个网络设备
    for (pcap_if_t* dev = all_devs; dev != nullptr; dev = dev->next) {
        std::cerr << "Device: " << dev->name << " (" << dev->description << ")" << std::endl;
    }
    //释放上面用完的设备列表内存，防止内存泄露
    pcap_freealldevs(all_devs);

    const char* dev = "\\Device\\NPF_{9E9D29C8-B719-4130-A129-20BC1F5ABBDF} ";
    if (!dev) {
        std::cerr << "Device error: " << errbuf << std::endl;
        return 1;
    }

    //打开设备进行抓包（实时捕获）
    pcap_t* handle = pcap_open_live(dev, 1024 * 1024, 1, 5000, errbuf);
    if (!handle) {
        std::cerr << "Open error: " << errbuf << std::endl;
        return 2;
    }

    struct bpf_program fp;  //用于定义过滤规则，比如只捕获 TCP 包、特定端口、特定 IP 等
    if (pcap_compile(handle, &fp, "", 0, 0) == -1) {
        std::cerr << "Filter error: " << pcap_geterr(handle) << std::endl;
        return 3;
    }
    if (pcap_setfilter(handle, &fp) == -1) {
        std::cerr << "Set filter error: " << pcap_geterr(handle) << std::endl;
        return 4;
    }

    log_file.open("traffic.csv");
    log_file << "时间,源MAC,源IP,目标MAC,目标IP,帧长度\n";
    last_stat_time = time(nullptr);

    struct pcap_pkthdr* header;
    const u_char* packet;

    std::cerr << "begin!" << std::endl;
    //主循环：不断抓取数据包
    while (true) {
        int res = pcap_next_ex(handle, &header, &packet);
        std::cerr << "get one!" << ' ' << res << " ";
        if (res == 1) {
            std::cerr << header << ' ' << packet << std::endl;
            process_packet(header, packet);
        }
        else if (res == -1) {
            std::cerr << "Read error: " << pcap_geterr(handle) << std::endl;
            break;
        }

        if (time(nullptr) - last_stat_time >= 60) {
            output_stats();
            last_stat_time = time(nullptr);
        }
        if (_kbhit() && _getch() == 'q') {
            std::cout << "User requested stop." << std::endl;
            break;
        }
    }

    pcap_close(handle);
    log_file.close();
    return 0;
}